from concurrent.futures import ThreadPoolExecutor
from fastapi import FastAPI, File, Form
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from uvicorn import Config, Server
from fastapi.middleware.cors import CORSMiddleware
import threading
from subsystem import android, android_json, analyse_features_android, android_convert, android_predict
from subsystem import pe, pe_json, analyse_features_pe, pe_predict
from subsystem import elf, elf_json, analyse_features_elf, elf_predict
from loguru import logger
import sys
import json

logger.remove()  # 这行很关键，先删除logger自动产生的handler，不然会出现重复输出的问题
# logger.add(sys.stderr, level='INFO')  # 只输出警告以上的日志
logger.add(sys.stderr, level='WARNING')  # 只输出警告以上的日志

# 创建线程池
pool = ThreadPoolExecutor(max_workers=10)


class Response(BaseModel):
    code: int = 200
    message: str = ""
    data: dict = {}


class Keylist(BaseModel):
    keylist: list = []


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/test")
def test():
    return Response()


running_list_a = []
analysing_lsit_a = []
running_list_p = []
analysing_lsit_p = []
finish_list_a = []
finish_list_p = []
finish_dict_a = {}
finish_dict_p = {}
lock_a = threading.Lock()
lock_p = threading.Lock()


@app.post("/file/upload")
def upload_file(file: bytes = File(...),
                key: str = Form(...),
                flag: str = Form(...),
                type: str = Form(...)):
    if flag == "apk":
        target_lock = lock_a
        target_list = running_list_a
        target_func = android_detect
    elif flag == "pe":
        target_lock = lock_p
        target_list = running_list_p
        target_func = pe_detect
    else:  # elf 或 不支持类型(将会返回错误)
        target_lock = lock_p
        target_list = running_list_p
        target_func = elf_detect
    # 覆盖target_func
    if type == "feature":
        target_func = json_detect
    with target_lock:
        target_list.append(key)
    pool.submit(target_func, key, file)
    # t = threading.Thread(target=target_func, args=[key, file])
    # t.start()
    return Response()


def json_detect(key, file):
    data = json.loads(file)
    if (data["type"] == "apk"):
        android_json_detect(key, file)
    elif (data["type"] == "pe"):
        pe_json_detect(key, file)
    elif (data['type'] == "elf"):
        elf_json_detect(key, file)


@app.post("/file/analyse")
def analyse(file: bytes = File(...),
            key: str = Form(...),
            flag: str = Form(...)):
    if flag == "apk":
        target_lock = lock_a
        target_list = analysing_lsit_a
        target_func = analyse_android
    elif flag == "pe":
        target_lock = lock_p
        target_list = analysing_lsit_p
        target_func = analyse_pe
    elif flag == "elf":
        target_lock = lock_p
        target_list = analysing_lsit_p
        target_func = analyse_elf
    with target_lock:
        target_list.append(key)
    pool.submit(target_func, key, file)
    # t = threading.Thread(target=target_func, args=[key, file])
    # t.start()
    return Response()


@app.post("/android/convert_features")
def convert_features(file: bytes = File(...), key: str = Form(...)):
    flag = android_convert(file, key)
    if flag:
        return Response()
    else:
        return Response(code=400)


@app.post("/android/detect_vec")
def android_detect_vec(file: bytes = File(...)):
    res = android_predict(file)
    if res:
        return Response(data={"pred": res})
    else:
        return Response(code=400)


@app.post("/pe/detect_vec")
def pe_detect_vec(file: bytes = File(...)):
    data = json.loads(file)
    if (data["type"] == "pe"):
        res = pe_predict(data)
    elif (data["type"] == "elf"):
        res = elf_predict(data)
    if res:
        return Response(data={"pred": res})
    else:
        return Response(code=400)


def analyse_android(key, file):
    flag = analyse_features_android(file, key)
    if flag:
        state = "done"
    else:
        state = "error"
    with lock_a:
        analysing_lsit_a.remove(key)
        finish_list_a.append(key)
        finish_dict_a[key] = state


def analyse_pe(key, file):
    flag = analyse_features_pe(file, key)
    if flag:
        state = "done"
    else:
        state = "error"
    with lock_p:
        analysing_lsit_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = state


def analyse_elf(key, file):
    flag = analyse_features_elf(file, key)
    if flag:
        state = "done"
    else:
        state = "error"
    with lock_p:
        analysing_lsit_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = state


@app.post("/state/android")
def query_android_state(keylist: Keylist):
    # print(keylist)
    keys = keylist.keylist
    logger.info("query android state")
    with lock_a:
        running = list(set(running_list_a) & set(keys))
        finish = list(set(finish_list_a) & set(keys))
        analysing = list(set(analysing_lsit_a) & set(keys))
        # logger.info("running:"+str(running_list_a))
        # logger.info("finish:"+str(finish_list_a))
        # logger.info("key:"+str(keys))
    state = {}
    for item in running:
        state[item] = "running"
    for item in analysing:
        state[item] = "analysing"
    for item in finish:
        state[item] = finish_dict_a[item]
        del finish_dict_a[item]
        finish_list_a.remove(item)
    return Response(data={"state": state})


@app.post("/state/pe")
def query_pe_state(keylist: Keylist):
    # print(keylist)
    keys = keylist.keylist
    logger.info("query pe state")
    with lock_p:
        running = list(set(running_list_p) & set(keys))
        finish = list(set(finish_list_p) & set(keys))
        analysing = list(set(analysing_lsit_p) & set(keys))
    state = {}
    for item in running:
        state[item] = "running"
    for item in analysing:
        state[item] = "analysing"
    for item in finish:
        state[item] = finish_dict_p[item]
        del finish_dict_p[item]
        finish_list_p.remove(item)
    return Response(data={"state": state})


app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/download/android/{file_key}")
def get_android_file(file_key):
    logger.info(f"download:{file_key}")
    return FileResponse(f"static/files/{file_key}.json")


@app.get("/download/pe/{file_key}")
def get_pe_file(file_key):
    logger.info(f"download:{file_key}")
    return FileResponse(f"static/files/{file_key}.json")


def android_detect(key, file):
    res = android(file, key)
    if res == "error":
        logger.error("Android File:" + key + " analyse error")
    with lock_a:
        running_list_a.remove(key)
        finish_list_a.append(key)
        finish_dict_a[key] = res


def android_json_detect(key, file):
    res = android_json(file, key)
    if res == "error":
        logger.error("Android Json File:" + key + " analyse error")
    with lock_a:
        running_list_a.remove(key)
        finish_list_a.append(key)
        finish_dict_a[key] = res


def pe_detect(key, file):
    res = pe(file, key)
    if res == "error":
        logger.error("PE File:" + key + " analyse error")
    with lock_p:
        running_list_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = res


def pe_json_detect(key, file):
    res = pe_json(file, key)
    if res == "error":
        logger.error("PE Json File:" + key + " analyse error")
    with lock_p:
        running_list_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = res


def elf_detect(key, file):
    res = elf(file, key)
    if res == "error":
        logger.error("ELF File:" + key + " analyse error")
    with lock_p:
        running_list_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = res


def elf_json_detect(key, file):
    res = elf_json(file, key)
    if res == "error":
        logger.error("ELF Json File:" + key + " analyse error")
    with lock_p:
        running_list_p.remove(key)
        finish_list_p.append(key)
        finish_dict_p[key] = res


if __name__ == '__main__':
    config = Config(app=app, reload=True, port=8888)
    server = Server(config=config)
    server.run()
